from flask import Flask, send_file, Response, render_template
import cv2
import io
import numpy as np
from ultralytics import YOLO
import os
import torch
app = Flask(__name__)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

current_directory = os.getcwd()
model_path = os.path.join(current_directory, 'yolo11n.pt')

# model = YOLO(model_path)
model = YOLO(model_path).to(device)
cap = cv2.VideoCapture(0)

# def draw_labels(frame, preds):
#     for pred in preds:
#         x1, y1, x2, y2, conf, cls = pred
#         label = f"{cls} {conf:.2f}"
#         color = (0, 255, 0)
#         cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.45, color, 2)
#         cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
#     return frame

def get_frame():
    # 假设 cap 已经被正确初始化并可以读取视频帧
    ret, frame = cap.read()
    if not ret:
        return None

        # 假设 model 是已经加载并可以调用的模型
    results = model(frame, conf=0.25, iou=0.45)

    # 如果模型返回的结果为空，则直接返回原始帧的字节流
    if not results:
        ret, jpeg = cv2.imencode('.jpg', frame)
        return jpeg.tobytes()

        # 创建一个用于绘制的副本，避免直接修改原始帧
    annotated_frame = frame.copy()

    # 遍历所有检测结果
    for result in results:
        # 每个 result 对象都有 plot 方法，用于在图像上绘制检测框
        annotated_frame = result.plot()

    # 将图像转换为字节流
    ret, jpeg = cv2.imencode('.jpg', annotated_frame)
    return jpeg.tobytes()

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/video_feed')
def video_feed():
    frame = get_frame()
    if frame is None:
        return Response(status=404)
    return Response(frame, mimetype='image/jpeg')


if __name__ == '__main__':
    app.run(host='0.0.0.0', debug=True, threaded=True)